#!/usr/bin/env python
# encoding: utf-8
"""
检查 LoRA 冻结是否生效的验证脚本
运行此脚本可以确认基础模型参数是否被正确冻结
"""

import torch
import torch.nn as nn
import sys
import os

sys.path.append("src")

def check_lora_freeze(model):
    """
    检查模型中 LoRA 参数的冻结状态
    
    Returns:
        dict: 包含冻结状态统计信息
    """
    stats = {
        "total_params": 0,
        "trainable_params": 0,
        "frozen_params": 0,
        "lora_params": 0,
        "base_params": 0,
        "trainable_lora_params": 0,
        "frozen_base_params": 0
    }
    
    # 检查 LoRALinear 模块
    from src.lora.lucaone_lora_adapter import LoRALinear
    
    for name, module in model.named_modules():
        if isinstance(module, LoRALinear):
            # LoRA 参数
            if hasattr(module, 'lora_A'):
                stats["lora_params"] += module.lora_A.numel()
                if module.lora_A.requires_grad:
                    stats["trainable_lora_params"] += module.lora_A.numel()
                    stats["trainable_params"] += module.lora_A.numel()
                else:
                    stats["frozen_params"] += module.lora_A.numel()
            
            if hasattr(module, 'lora_B'):
                stats["lora_params"] += module.lora_B.numel()
                if module.lora_B.requires_grad:
                    stats["trainable_lora_params"] += module.lora_B.numel()
                    stats["trainable_params"] += module.lora_B.numel()
                else:
                    stats["frozen_params"] += module.lora_B.numel()
            
            # 基础层参数（应该被冻结）
            for param_name, param in module.base_layer.named_parameters():
                stats["base_params"] += param.numel()
                if param.requires_grad:
                    print(f"⚠️  警告: 基础层参数 {name}.base_layer.{param_name} 未被冻结!")
                    stats["trainable_params"] += param.numel()
                else:
                    stats["frozen_base_params"] += param.numel()
                    stats["frozen_params"] += param.numel()
    
    # 检查其他参数（非 LoRA 模块）
    for name, param in model.named_parameters():
        stats["total_params"] += param.numel()
        
        # 跳过已经在 LoRA 模块中统计的参数
        if "lora_A" in name or "lora_B" in name:
            continue
        if ".base_layer." in name:
            continue
        
        if param.requires_grad:
            stats["trainable_params"] += param.numel()
        else:
            stats["frozen_params"] += param.numel()
    
    return stats


def print_freeze_stats(stats):
    """打印冻结状态统计"""
    print("\n" + "=" * 60)
    print("LoRA 冻结状态检查")
    print("=" * 60)
    print(f"总参数数量: {stats['total_params']:,}")
    print(f"可训练参数: {stats['trainable_params']:,} ({100*stats['trainable_params']/stats['total_params']:.2f}%)")
    print(f"冻结参数: {stats['frozen_params']:,} ({100*stats['frozen_params']/stats['total_params']:.2f}%)")
    print("\n" + "-" * 60)
    print("LoRA 参数:")
    print(f"  LoRA 参数总数: {stats['lora_params']:,}")
    print(f"  可训练 LoRA 参数: {stats['trainable_lora_params']:,} ({100*stats['trainable_lora_params']/stats['lora_params']:.2f}%)")
    print("\n基础层参数:")
    print(f"  基础层参数总数: {stats['base_params']:,}")
    print(f"  冻结的基础层参数: {stats['frozen_base_params']:,} ({100*stats['frozen_base_params']/stats['base_params']:.2f}%)")
    print("=" * 60)
    
    # 检查是否正常
    if stats['trainable_params'] / stats['total_params'] > 0.1:
        print("\n⚠️  警告: 可训练参数占比过高 (>10%)!")
        print("   可能原因:")
        print("   1. 基础模型参数未被正确冻结")
        print("   2. 分类头或其他模块参数未被冻结")
        print("   建议: 检查 freeze_base_model() 是否正确执行")
    else:
        print("\n✓ LoRA 冻结状态正常 (可训练参数 < 10%)")


if __name__ == "__main__":
    print("LoRA 冻结验证脚本")
    print("=" * 60)
    print("\n使用方法:")
    print("1. 在训练脚本中添加此检查:")
    print("   from check_lora_freeze import check_lora_freeze, print_freeze_stats")
    print("   stats = check_lora_freeze(model)")
    print("   print_freeze_stats(stats)")
    print("\n2. 或者直接运行此脚本（需要先加载模型）")
    print("\n注意: 此脚本需要在模型加载并应用 LoRA 后运行")

